package cc.git.liuyan.customeraiagent.core.vectordatabase;

import cc.git.liuyan.customeraiagent.core.embeddingmodel.BigDecimalVec;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.collection.response.ListCollectionsResp;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.ranker.RRFRanker;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.extern.slf4j.Slf4j;

import java.math.BigDecimal;
import java.util.*;

@Slf4j
public class MilvusVectorDataBase implements VectorDataBase {
    private String uri;
    private String token;
    private MilvusClientV2 milvusClientV2;
    private final String knowledgeCollectionName = "customeraiagent_knowledge_vector_collection";
    private final String text2SqlCollectionName = "customeraiagent_text2sql_vector_collection";

    public MilvusVectorDataBase(String uri, String token) {
        this.uri = uri;
        this.token = token;
        ConnectConfig connectConfig = ConnectConfig.builder().uri(this.uri).token(this.token).build();
        milvusClientV2 = new MilvusClientV2(connectConfig);
    }

    @Override
    public void createKnowledgeVectorCollection(int denseDimension, int sparseDimension) {
        ListCollectionsResp resp = milvusClientV2.listCollections();
        if (resp != null && resp.getCollectionNames() != null && resp.getCollectionNames().contains(knowledgeCollectionName)) {
            //已创建
        } else {
            CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema();
            schema.setEnableDynamicField(true);
            {
                schema.addField(AddFieldReq.builder()
                        .fieldName("Auto_id")
                        .dataType(DataType.Int64)
                        .isPrimaryKey(true)
                        .autoID(true)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("denseVectors")
                        .dataType(DataType.FloatVector)
                        .dimension(denseDimension)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("sparseVectors")
                        .dataType(DataType.FloatVector)
                        .dimension(sparseDimension)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("knowledgeId")
                        .dataType(DataType.Int64)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("segmentNum")
                        .dataType(DataType.Int64)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("segmentContent")
                        .dataType(DataType.VarChar)
                        .maxLength(2000)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("score")
                        .dataType(DataType.Int64)
                        .build());
            }

            List<IndexParam> indexParams = new ArrayList<>();
            {
                Map<String, Object> denseParams = new HashMap<>();
                denseParams.put("nlist", 128);
                IndexParam denseVectorIndex = IndexParam.builder()
                        .fieldName("denseVectors")
                        .indexName("denseVectors_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .metricType(IndexParam.MetricType.IP)
                        .extraParams(denseParams)
                        .build();

                Map<String, Object> sparseParams = new HashMap<>();
                sparseParams.put("drop_ratio_build", 0.2);
                IndexParam sparseVectorIndex = IndexParam.builder()
                        .fieldName("sparseVectors")
                        .indexName("sparseVectors_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .metricType(IndexParam.MetricType.IP)
                        .extraParams(sparseParams)
                        .build();

                IndexParam knowledgeIdIndex = IndexParam.builder()
                        .fieldName("knowledgeId")
                        .indexName("knowledgeId_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .build();

                indexParams.add(denseVectorIndex);
                indexParams.add(sparseVectorIndex);
                indexParams.add(knowledgeIdIndex);
            }

            CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
                    .collectionName(knowledgeCollectionName)
                    .collectionSchema(schema)
                    .indexParams(indexParams)
                    .enableDynamicField(true)
                    .build();

            milvusClientV2.createCollection(createCollectionReq);

            GetLoadStateReq getLoadStateReq = GetLoadStateReq.builder()
                    .collectionName(knowledgeCollectionName)
                    .build();

            try {
                Boolean loaded = false;
                int maxChechNum = 10;
                while (!loaded && maxChechNum-- > 0) {
                    loaded = milvusClientV2.getLoadState(getLoadStateReq);
                    Thread.sleep(1000L);
                }
                if (loaded) {
                    return;
                } else {
                    throw new Exception("Failed to get load state for knowledge collection");
                }
            } catch (Exception ex) {
                log.error("createKnowledgeVectorCollection  getLoadState  error");
                log.error(ex.getMessage(), ex);
            }
        }
    }

    @Override
    public void deleteKnowledgeTrain(Long knowledgeId) {
        String filterExpr = "knowledgeId == " + knowledgeId;
        DeleteReq deleteReq = DeleteReq.builder().collectionName(knowledgeCollectionName).filter(filterExpr).build();
        milvusClientV2.delete(deleteReq);
    }

    @Override
    public void trainKowledge(KnowledgeVectorData knowledgeVectorData) {
        List<JsonObject> data = new ArrayList<>();
        JsonArray denseJsonArray = new JsonArray();
        for (BigDecimal value : knowledgeVectorData.getDenseVectors()) {
            denseJsonArray.add(value);
        }
        JsonArray sparseJsonArray = new JsonArray();
        for (BigDecimal value : knowledgeVectorData.getSparseVectors()) {
            sparseJsonArray.add(value);
        }
        JsonObject jsonObject = new JsonObject();
        jsonObject.add("denseVectors", denseJsonArray);
        jsonObject.add("sparseVectors", sparseJsonArray);
        jsonObject.addProperty("knowledgeId", knowledgeVectorData.getKnowledgeId());
        jsonObject.addProperty("segmentNum", knowledgeVectorData.getSegmentNum());
        jsonObject.addProperty("segmentContent", knowledgeVectorData.getSegmentContent());
        jsonObject.addProperty("score", knowledgeVectorData.getScore());
        data.add(jsonObject);
        InsertReq insertReq = InsertReq.builder().collectionName(knowledgeCollectionName).data(data).build();
        milvusClientV2.insert(insertReq);
    }

    @Override
    public List<String> relationKnowledgeSearch(List<Long> knowledgeIds, List<BigDecimal> denseVectors, List<BigDecimal> sparseVectors) {
        List<BaseVector> queryDenseVectors = Collections.singletonList(new BigDecimalVec(denseVectors));
        List<BaseVector> querySparseVectors = Collections.singletonList(new BigDecimalVec(sparseVectors));
        StringJoiner knowlwdgeIdJoiner = new StringJoiner(", ");
        knowledgeIds.forEach(knowledgeId -> knowlwdgeIdJoiner.add(knowledgeId.toString()));
        String filterExpr = "knowledgeId in [ " + knowlwdgeIdJoiner.toString() + " ]";
        List<AnnSearchReq> searchRequests = new ArrayList<>();
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("denseVectors")
                .vectors(queryDenseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{\"nprobe\": 10}")
                .expr(filterExpr)
                .topK(10)
                .build());
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("sparseVectors")
                .vectors(querySparseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{\"drop_ratio_build\": 0.2}")
                .expr(filterExpr)
                .topK(10)
                .build());
        HybridSearchReq hybridSearchReq = HybridSearchReq.builder()
                .collectionName(knowledgeCollectionName)
                .searchRequests(searchRequests)
                .ranker(new RRFRanker(60))
                .topK(10)
                .outFields(Collections.singletonList("segmentContent"))
                .build();
        SearchResp searchResp = milvusClientV2.hybridSearch(hybridSearchReq);
        List<String> relationSegmentContents = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        for (List<SearchResp.SearchResult> results : searchResults) {
            for (SearchResp.SearchResult result : results) {
                String segmentContent = result.getEntity().get("segmentContent").toString();
                relationSegmentContents.add(segmentContent);
            }
        }
        return relationSegmentContents;
    }

    @Override
    public void createText2SqlVectorCollection(int denseDimension, int sparseDimension) {
        ListCollectionsResp resp = milvusClientV2.listCollections();
        if (resp != null && resp.getCollectionNames() != null && resp.getCollectionNames().contains(text2SqlCollectionName)) {
            //已创建
        } else {
            CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema();
            schema.setEnableDynamicField(true);
            {
                schema.addField(AddFieldReq.builder()
                        .fieldName("Auto_id")
                        .dataType(DataType.Int64)
                        .isPrimaryKey(true)
                        .autoID(true)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("denseVectors")
                        .dataType(DataType.FloatVector)
                        .dimension(denseDimension)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("sparseVectors")
                        .dataType(DataType.FloatVector)
                        .dimension(sparseDimension)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("knowledgeId")
                        .dataType(DataType.Int64)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("type")
                        .dataType(DataType.VarChar)
                        .description("ddl/remark")
                        .maxLength(200)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("schema")
                        .dataType(DataType.VarChar)
                        .maxLength(200)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("tableName")
                        .dataType(DataType.VarChar)
                        .maxLength(200)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("ddl")
                        .dataType(DataType.VarChar)
                        .maxLength(10000)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("remarkId")
                        .dataType(DataType.Int64)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("remark")
                        .dataType(DataType.VarChar)
                        .maxLength(2000)
                        .build());
                schema.addField(AddFieldReq.builder()
                        .fieldName("score")
                        .dataType(DataType.Int64)
                        .build());
            }

            List<IndexParam> indexParams = new ArrayList<>();
            {
                Map<String, Object> denseParams = new HashMap<>();
                denseParams.put("nlist", 128);
                IndexParam denseVectorIndex = IndexParam.builder()
                        .fieldName("denseVectors")
                        .indexName("denseVectors_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .metricType(IndexParam.MetricType.IP)
                        .extraParams(denseParams)
                        .build();

                Map<String, Object> sparseParams = new HashMap<>();
                sparseParams.put("drop_ratio_build", 0.2);
                IndexParam sparseVectorIndex = IndexParam.builder()
                        .fieldName("sparseVectors")
                        .indexName("sparseVectors_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .metricType(IndexParam.MetricType.IP)
                        .extraParams(sparseParams)
                        .build();

                IndexParam knowledgeIdIndex = IndexParam.builder()
                        .fieldName("knowledgeId")
                        .indexName("knowledgeId_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .build();

                IndexParam typeIdIndex = IndexParam.builder()
                        .fieldName("type")
                        .indexName("type_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .build();

                IndexParam schemaIndex = IndexParam.builder()
                        .fieldName("schema")
                        .indexName("schema_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .build();

                IndexParam tableNameIndex = IndexParam.builder()
                        .fieldName("tableName")
                        .indexName("tableName_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .build();

                IndexParam remarkIdIndex = IndexParam.builder()
                        .fieldName("remarkId")
                        .indexName("remarkId_index")
                        .indexType(IndexParam.IndexType.AUTOINDEX)
                        .build();

                indexParams.add(denseVectorIndex);
                indexParams.add(sparseVectorIndex);
                indexParams.add(knowledgeIdIndex);
                indexParams.add(typeIdIndex);
                indexParams.add(schemaIndex);
                indexParams.add(tableNameIndex);
                indexParams.add(remarkIdIndex);
            }

            CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
                    .collectionName(text2SqlCollectionName)
                    .collectionSchema(schema)
                    .indexParams(indexParams)
                    .enableDynamicField(true)
                    .build();

            milvusClientV2.createCollection(createCollectionReq);

            GetLoadStateReq getLoadStateReq = GetLoadStateReq.builder()
                    .collectionName(text2SqlCollectionName)
                    .build();

            try {
                Boolean loaded = false;
                int maxChechNum = 10;
                while (!loaded && maxChechNum-- > 0) {
                    loaded = milvusClientV2.getLoadState(getLoadStateReq);
                    Thread.sleep(1000L);
                }
                if (loaded) {
                    return;
                } else {
                    throw new Exception("failed to get load state for text2sql collection");
                }
            } catch (Exception ex) {
                log.error("createText2SqlVectorCollection  getLoadState  error");
                log.error(ex.getMessage(), ex);
            }
        }
    }

    @Override
    public void trainText2Sql(MysqlSchemaVectorData mysqlSchemaVectorData) {
        List<JsonObject> data = new ArrayList<>();
        JsonArray denseJsonArray = new JsonArray();
        for (BigDecimal value : mysqlSchemaVectorData.getDenseVectors()) {
            denseJsonArray.add(value);
        }
        JsonArray sparseJsonArray = new JsonArray();
        for (BigDecimal value : mysqlSchemaVectorData.getSparseVectors()) {
            sparseJsonArray.add(value);
        }
        JsonObject jsonObject = new JsonObject();
        jsonObject.add("denseVectors", denseJsonArray);
        jsonObject.add("sparseVectors", sparseJsonArray);
        jsonObject.addProperty("knowledgeId", mysqlSchemaVectorData.getKnowledgeId());
        jsonObject.addProperty("type", mysqlSchemaVectorData.getType());
        jsonObject.addProperty("schema", mysqlSchemaVectorData.getSchema());
        jsonObject.addProperty("tableName", mysqlSchemaVectorData.getTableName());
        jsonObject.addProperty("ddl", mysqlSchemaVectorData.getDdl());
        jsonObject.addProperty("remarkId", mysqlSchemaVectorData.getRemarkId());
        jsonObject.addProperty("remark", mysqlSchemaVectorData.getRemark());
        jsonObject.addProperty("score", mysqlSchemaVectorData.getScore());
        data.add(jsonObject);
        InsertReq insertReq = InsertReq.builder().collectionName(text2SqlCollectionName).data(data).build();
        milvusClientV2.insert(insertReq);
    }

    @Override
    public void deleteText2SqlTrainRemark(Long knowledgeId, Long remarkId) {
        String filterExpr = String.format("knowledgeId == %s AND remarkId == %s", knowledgeId, remarkId);
        DeleteReq deleteReq = DeleteReq.builder().collectionName(text2SqlCollectionName).filter(filterExpr).build();
        milvusClientV2.delete(deleteReq);
    }

    @Override
    public void deleteText2SqlTrainDdl(Long knowledgeId, String ddlSchema, String ddlTable) {
        String filterExpr = String.format("knowledgeId == %s AND schema == \"%s\" AND tableName == \"%s\"", knowledgeId, ddlSchema, ddlTable);
        DeleteReq deleteReq = DeleteReq.builder().collectionName(text2SqlCollectionName).filter(filterExpr).build();
        milvusClientV2.delete(deleteReq);
    }

    @Override
    public List<String> relationSchemaSearch(List<BigDecimal> denseVectors, List<BigDecimal> sparseVectors) {
        List<BaseVector> queryDenseVectors = Collections.singletonList(new BigDecimalVec(denseVectors));
        List<BaseVector> querySparseVectors = Collections.singletonList(new BigDecimalVec(sparseVectors));
        String filterExpr = String.format("type == \"%s\"", "ddl");
        List<AnnSearchReq> searchRequests = new ArrayList<>();
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("denseVectors")
                .vectors(queryDenseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{\"nprobe\": 10}")
                .expr(filterExpr)
                .topK(10)
                .build());
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("sparseVectors")
                .vectors(querySparseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{\"drop_ratio_build\": 0.2}")
                .expr(filterExpr)
                .topK(10)
                .build());
        HybridSearchReq hybridSearchReq = HybridSearchReq.builder()
                .collectionName(text2SqlCollectionName)
                .searchRequests(searchRequests)
                .ranker(new RRFRanker(60))
                .topK(10)
                .outFields(Collections.singletonList("ddl"))
                .build();
        SearchResp searchResp = milvusClientV2.hybridSearch(hybridSearchReq);
        List<String> relationDdls = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        for (List<SearchResp.SearchResult> results : searchResults) {
            for (SearchResp.SearchResult result : results) {
                String ddl = result.getEntity().get("ddl").toString();
                relationDdls.add(ddl);
            }
        }
        return relationDdls;
    }

    @Override
    public List<String> relationRemarkSearch(List<BigDecimal> denseVectors, List<BigDecimal> sparseVectors) {
        List<BaseVector> queryDenseVectors = Collections.singletonList(new BigDecimalVec(denseVectors));
        List<BaseVector> querySparseVectors = Collections.singletonList(new BigDecimalVec(sparseVectors));
        String filterExpr = String.format("type == \"%s\"", "remark");
        List<AnnSearchReq> searchRequests = new ArrayList<>();
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("denseVectors")
                .vectors(queryDenseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{\"nprobe\": 10}")
                .expr(filterExpr)
                .topK(10)
                .build());
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("sparseVectors")
                .vectors(querySparseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{\"drop_ratio_build\": 0.2}")
                .expr(filterExpr)
                .topK(10)
                .build());
        HybridSearchReq hybridSearchReq = HybridSearchReq.builder()
                .collectionName(text2SqlCollectionName)
                .searchRequests(searchRequests)
                .ranker(new RRFRanker(60))
                .topK(10)
                .outFields(Collections.singletonList("remark"))
                .build();
        SearchResp searchResp = milvusClientV2.hybridSearch(hybridSearchReq);
        List<String> relationRemarks = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        for (List<SearchResp.SearchResult> results : searchResults) {
            for (SearchResp.SearchResult result : results) {
                String remark = result.getEntity().get("remark").toString();
                relationRemarks.add(remark);
            }
        }
        return relationRemarks;
    }
}
